#!/usr/bin/env python


import rospy
from std_msgs.msg import String
import ros_numpy
import numpy as np
from sensor_msgs.msg import PointCloud2,PointCloud
from geometry_msgs.msg import PoseStamped,Pose,PoseWithCovarianceStamped,TransformStamped
from ros_numpy import numpify,msgify
from nav_msgs.msg import Odometry
from apriltags2_ros.msg import AprilTagDetectionArray,AprilTagDetection

import tf2_ros
from tf.transformations import euler_from_quaternion, quaternion_from_euler,quaternion_matrix,euler_matrix


def pose_to_tf(pose,parent_id,child_id):
    P=pose
    
    position = P.position
    orientation = P.orientation
    
    euler_angles = euler_from_quaternion([orientation.x, orientation.y, orientation.z, orientation.w])
    roll, pitch, yaw = euler_angles

    # Create a transform broadcaster
    br = tf2_ros.TransformBroadcaster()

    # Create a transform message
    t = TransformStamped()
    t.header.stamp = rospy.Time.now()
    t.header.frame_id = parent_id  # name of the parent frame
    t.child_frame_id = child_id  # name of the child frame

    # Set the translation
    t.transform.translation.x = position.x
    t.transform.translation.y = position.y
    t.transform.translation.z = position.z

    # Set the rotation
    q = quaternion_from_euler(roll, pitch, yaw)
    t.transform.rotation.x = q[0]
    t.transform.rotation.y = q[1]
    t.transform.rotation.z = q[2]
    t.transform.rotation.w = q[3]

    # Publish the transform
    br.sendTransform(t)
    


def pc_msg_2_np(pc_msg):

    points = []
    for point in pc_msg.points:
        x = point.x
        y = point.y
        z = point.z
        points.append([x, y, z])

    pc_np=np.array(points).T
    
    return pc_np


def pc2_msg_to_np(pc_msg):
    pc_np=numpify(pc_msg)
    pc_np= np.stack([pc_np[f] for f in ('x', 'y', 'z')])

    return pc_np


def np_to_pc2_msg(pc_np):
    data=np.zeros(pc_np.shape[1],
    dtype=[('x',np.float32), ('y',np.float32), ('z',np.float32)])
    data['x']=pc_np[0,]
    data['y']=pc_np[1,]
    data['z']=pc_np[2,]

    pc_msg=msgify(PointCloud2,data)
    pc_msg.is_dense=False


    
    return pc_msg


class Agent:
    
    def __init__(self):
        
        rospy.init_node('active_slot_node')
        
        self.gimble_angle=0
        
        self.T_C_W=np.eye(4) #Camera in World frame
        self.T_W_C=np.eye(4) 
        self.T_T_C=np.eye(4) #Target in Camera frame
        self.T_B_C=np.eye(4) #Camera in Body frame
        self.T_B_V=np.eye(4) #Body in Vrpn frame
        self.T_B_W=np.eye(4)
        
        self.pc_C=np.zeros([3,100])
        self.first_vins_pc_msg=True
        
        rospy.Subscriber('/vins_estimator/keyframe_pose',Odometry,self.kf_pose_cb)
        rospy.Subscriber('/vins_estimator/keyframe_point',PointCloud,self.kf_pc_cb)
        rospy.Subscriber('/tag_detections',AprilTagDetectionArray,self.tag_cb)
        rospy.Subscriber('/vrpn_client_node/uav_dy/pose',PoseStamped,self.vrpn_pose_cb)
        
        self.target_pose_pub=rospy.Publisher('/slot/target_pose',PoseWithCovarianceStamped,queue_size=10)
        self.pc2_C_pub=rospy.Publisher('/slot/pointcloud2',PointCloud2,queue_size=10)
        self.body_pose_pub=rospy.Publisher('/slot/body_pose',PoseStamped,queue_size=10)
    def vrpn_pose_cb(self,vrpn_msg):
        self.T_B_V=numpify(vrpn_msg.pose)
        T_V_B=np.linalg.inv(self.T_B_V)
        
        pose_to_tf(msgify(Pose,T_V_B),parent_id='body',child_id='vrpn')
        
                #calculate body pose
        self.T_B_C=euler_matrix(0.,-np.pi/2+self.gimble_angle,0.)
        self.T_B_W=self.T_C_W @self.T_B_C 
        body_pose_W_msg=PoseStamped()
        #body_pose_W_msg.header=camera_pose_msg.header
        body_pose_W_msg.header.frame_id="world"
        body_pose_W_msg.pose=msgify(Pose,self.T_B_W)
        self.body_pose_pub.publish(body_pose_W_msg)
        
        pose_to_tf(msgify(Pose,self.T_B_C),parent_id="dfi",child_id="body")
    
    def tag_cb(self,tag_msg):
        
        if(len(tag_msg.detections)>0):
        
            tag_detection=tag_msg.detections[0]
            #tag_detection=AprilTagDetection()
            
            target_pose_C=tag_detection.pose.pose.pose
            self.T_T_C=numpify(target_pose_C)
            
            target_pose_C_msg=PoseWithCovarianceStamped()
            target_pose_C_msg.header.frame_id='body'
            
            T_T_W=self.T_C_W @ self.T_T_C
            T_T_B=np.linalg.inv(self.T_B_C) @ self.T_T_C
            
            target_pose_C_msg.pose.pose=msgify(Pose,T_T_B)
            target_pose_C_msg.pose.pose.position.x+=1.3
            
            pose_to_tf(pose=target_pose_C_msg.pose.pose,parent_id="body",child_id="target")
 
            
            self.target_pose_pub.publish(target_pose_C_msg)
        

    def kf_pc_cb(self,pc_msg):
        
        points = []
        for point in pc_msg.points:
            x = point.x
            y = point.y
            z = point.z
            points.append([x, y, z,1.0])

        pc_W_aug=np.array(points).T
        self.pc_C_aug=self.T_W_C @ pc_W_aug
        self.pc_B_aug=np.linalg.inv(self.T_B_W) @ pc_W_aug
        
        pc_C=self.pc_C_aug[:3,:]
        pc_B=self.pc_B_aug[:3,:]
        
        pc_B_msg=np_to_pc2_msg(pc_B)
        pc_B_msg.header=pc_msg.header
        
        pc_B_msg.header.stamp=rospy.Time.now()
        pc_B_msg.header.frame_id="body"
        
        self.pc2_C_pub.publish(pc_B_msg)
        

    def kf_pose_cb(self,camera_pose_msg):
        
        if self.first_vins_pc_msg:
            #print(camera_pose_msg)
            self.first_vins_pc_msg=False
        
        P=Pose()
        P.position.x=camera_pose_msg.pose.pose.position.x
        P.position.y=camera_pose_msg.pose.pose.position.y
        P.position.z=camera_pose_msg.pose.pose.position.z

        P.orientation.x=camera_pose_msg.pose.pose.orientation.x
        P.orientation.y=camera_pose_msg.pose.pose.orientation.y
        P.orientation.z=camera_pose_msg.pose.pose.orientation.z
        P.orientation.w=camera_pose_msg.pose.pose.orientation.w


        self.T_C_W=numpify(P)
        self.T_W_C=np.linalg.inv(self.T_C_W)


        #publish tf
        pose_to_tf(P,parent_id="world",child_id="dfi")
        
    

        
        
if __name__ == '__main__':
    
    a=Agent()
    rospy.spin()


        
        